import torch
import matplotlib.pyplot as plt

# 构建x,y点图
x1 = torch.linspace(0, 1, 20)
w1 = torch.Tensor([2.])
w2 = torch.Tensor([-1.])
b1 = torch.Tensor([4.])
b2 = torch.Tensor([1.])
x2 = w1 * x1 + b1
y = x2 * w2 + b2

plt.plot(x1.detach().numpy(), y.detach().numpy(), 'ro')

# 训练
w1_predict = torch.Tensor([0.1])
w2_predict = torch.Tensor([0.1])
b1_predict = torch.Tensor([0.1])
b2_predict = torch.Tensor([0.1])
line, = plt.plot(x1, (w1_predict * x1 + b1_predict) * w2_predict + b2_predict, 'b--')

epochs = 100
for epoch in range(epochs):
    y_predict = (w1_predict * x1 + b1_predict) * w2_predict + b2_predict
    e = torch.mean((y - y_predict) ** 2)
    w1_predict -= torch.mean(-2 * x1 * w2 * (y - y_predict)) * 0.01
    w2_predict -= torch.mean(-2 * x2 * (y - y_predict)) * 0.01
    b1_predict -= torch.mean(-2 * w2 * (y - y_predict)) * 0.01
    b2_predict -= torch.mean(-2 * (y - y_predict)) * 0.01

    line.set_data(x1, y_predict)
    plt.pause(0.5)

# plt.show()
